# Copyright (c) 2000 SecureAuth Corporation.  All rights
# reserved.

# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:

# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.

# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in
#    the documentation and/or other materials provided with the
#    distribution.

# 3. The end-user documentation included with the redistribution,
#    if any, must include the following acknowledgment:
#       "This product includes software developed by
#        SecureAuth Corporation (https://www.secureauth.com/)."
#    Alternately, this acknowledgment may appear in the software itself,
#    if and wherever such third-party acknowledgments normally appear.

# 4. The names "Impacket", "SecureAuth Corporation" must
#    not be used to endorse or promote products derived from this
#    software without prior written permission. For written
#    permission, please contact oss@secureauth.com.

# 5. Products derived from this software may not be called "Impacket",
#    nor may "Impacket" appear in their name, without prior written
#    permission of SecureAuth Corporation.

# THIS SOFTWARE IS PROVIDED ``AS IS'' AND ANY EXPRESSED OR IMPLIED
# WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
# OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED.  IN NO EVENT SHALL THE APACHE SOFTWARE FOUNDATION OR
# ITS CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
# USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
# OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.


import logging
import time

from impacket.dcerpc.v5.dcom import wmi
from impacket.dcerpc.v5.dcomrt import DCOMConnection
from impacket.dcerpc.v5.dtypes import NULL
from impacket.smbconnection import SMBConnection

from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT

from .remote_shell import RemoteShell

logger = logging.getLogger(__name__)


# Adapted from https://github.com/SecureAuthCorp/impacket/blob/master/examples/wmiexec.py
# Used to get HKLM keys for restoring original DC password
class Wmiexec:
    OUTPUT_FILENAME = "__" + str(time.time())

    def __init__(
        self, ip, username, hashes, password="", domain="", share="ADMIN$", secrets_dir=None
    ):
        self.__ip = ip
        self.__username = username
        self.__password = password
        self.__domain = domain
        self.__lmhash, self.__nthash = hashes.split(":")
        self.__share = share
        self.__secrets_dir = secrets_dir
        self.shell = None

    def connect(self):
        self.smbConnection = SMBConnection(self.__ip, self.__ip, timeout=LONG_REQUEST_TIMEOUT)
        self.smbConnection.login(
            user=self.__username,
            password=self.__password,
            domain=self.__domain,
            lmhash=self.__lmhash,
            nthash=self.__nthash,
        )

        self.dcom = DCOMConnection(
            target=self.__ip,
            username=self.__username,
            password=self.__password,
            domain=self.__domain,
            lmhash=self.__lmhash,
            nthash=self.__nthash,
            oxidResolver=True,
        )

        try:
            iInterface = self.dcom.CoCreateInstanceEx(
                wmi.CLSID_WbemLevel1Login, wmi.IID_IWbemLevel1Login
            )
            iWbemLevel1Login = wmi.IWbemLevel1Login(iInterface)
            self.iWbemServices = iWbemLevel1Login.NTLMLogin("//./root/cimv2", NULL, NULL)
            iWbemLevel1Login.RemRelease()

        except (Exception, KeyboardInterrupt) as e:
            logger.error(str(e))
            self.smbConnection.logoff()
            self.dcom.disconnect()

    def get_remote_shell(self):
        self.connect()
        win32Process, _ = self.iWbemServices.GetObject("Win32_Process")
        self.shell = RemoteShell(
            self.__share, win32Process, self.smbConnection, self.OUTPUT_FILENAME, self.__secrets_dir
        )
        return self.shell

    def close(self):
        self.smbConnection.close()
        self.smbConnection = None
        self.dcom.disconnect()
        self.dcom = None
